#!/usr/bin/env python
import rospy
from std_msgs.msg import Bool
from geometry_msgs.msg import Pose, PoseStamped
from moveit_commander import MoveGroupCommander
import moveit_commander
import numpy as np
from tf.transformations import quaternion_matrix
import copy
import sys

# 全局变量
left_pose = None  # 接收到的目标位姿
last_valid_pose=None

start_move = False  # 是否开始任务
task_state = "stand_by"  # 任务状态：stand_by, move_to_paper, grasp, move_to_release, release, return_home

# MoveIt 组
move_group = None

# 手控制发布器
hand_control_publisher = None


def left_pose_callback(msg):
    """
    接收目标位姿的回调函数。
    """
    global left_pose, last_valid_pose

    # 检查当前 pose 是否有效（避免接收到无效位姿）
    if msg.pose.position.x == -1 or msg.pose.position.y == -1 or msg.pose.position.z == -1:
        if last_valid_pose:
            #rospy.loginfo("Invalid pose received, using last valid pose.")
            left_pose = copy.deepcopy(last_valid_pose)
        else:
            #rospy.logwarn("No valid pose available, skipping this callback.")
            return
    else:
        # 计算当前位置与上次有效位置之间的距离
        if last_valid_pose:
            current_position = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z])
            last_position = np.array([last_valid_pose.pose.position.x, last_valid_pose.pose.position.y, last_valid_pose.pose.position.z])
            distance = np.linalg.norm(current_position - last_position)

            # 如果距离大于 1cm，视为无效位姿，不进行规划
            if distance > 0.10:
                rospy.loginfo(f"Position change exceeds 1cm ({distance:.4f}m), skipping this pose.")
                return
            else:
                left_pose = copy.deepcopy(msg)
                last_valid_pose = copy.deepcopy(msg)  # 保存当前有效的 pose
        else:
            left_pose = copy.deepcopy(msg)
            last_valid_pose = copy.deepcopy(msg)  # 保存当前有效的 pose

    #rospy.loginfo("Received a pose, starting movement.")


def start_move_callback(msg):
    """
    接收开始命令的回调函数。
    """
    global start_move
    if msg.data:
        rospy.loginfo("Received start_move signal, starting task.")
        start_move = True
    else:
        rospy.loginfo("Received stop_move signal, stopping task.")
        start_move = False


def move_to_pose(pose, link_name):
    """
    移动手臂到指定位姿，使用分步规划策略。
    :param pose: Pose 消息，目标位姿
    :param link_name: 要移动的链接名称
    """
    global move_group
    print("Target pose:", pose)
    print("Current pose:", move_group.get_current_pose().pose)

    current_pose = move_group.get_current_pose().pose

    # 计算当前位置与目标位置之间的差距
    distance = np.linalg.norm(
        np.array([pose.pose.position.x, pose.pose.position.y, pose.pose.position.z]) -
        np.array([current_pose.position.x, current_pose.position.y, current_pose.position.z])
    )
    print(f"Distance to target pose: {distance}")

    # 设定一个阈值，如果目标位姿与当前位置的距离大于这个值，则进行分步规划
    step_size = 0.5  # 每一步的距离，单位米
    if distance > step_size:
        num_steps = int(np.ceil(distance / step_size))  # 计算需要的步数

        # 按照分步的策略逐步移动
        for i in range(num_steps):
            # 计算每一步的目标位置
            interpolation_factor = (i + 1) / num_steps  # 计算每一步的插值因子
            intermediate_pose = Pose()
            intermediate_pose.position.x = current_pose.position.x + (
                pose.pose.position.x - current_pose.position.x
            ) * interpolation_factor
            intermediate_pose.position.y = current_pose.position.y + (
                pose.pose.position.y - current_pose.position.y
            ) * interpolation_factor
            intermediate_pose.position.z = current_pose.position.z + (
                pose.pose.position.z - current_pose.position.z
            ) * interpolation_factor

            # 使用当前的位姿插值更新目标姿态
            intermediate_pose.orientation = pose.pose.orientation  # 保持目标的旋转信息不变

            move_group.set_pose_target(intermediate_pose)
            plan = move_group.plan()

            if plan[0]:
                success = move_group.execute(plan[1], wait=True)
                if not success:
                    rospy.logwarn(f"Failed to move {link_name} to intermediate pose at step {i + 1}.")
                    # 如果执行失败，直接尝试备选姿态
                    if try_alternative_orientations(intermediate_pose, link_name):
                        continue  # 如果备选姿态成功，继续下一步
                    else:
                        return False  # 如果备选姿态也失败，终止任务
            else:
                rospy.logwarn(f"Failed to generate plan for {link_name} at step {i + 1}.")
                # 如果规划失败，直接尝试备选姿态
                if try_alternative_orientations(intermediate_pose, link_name):
                    continue  # 如果备选姿态成功，继续下一步
                else:
                    return False  # 如果备选姿态也失败，终止任务

        rospy.loginfo(f"Successfully moved {link_name} to the target pose in {num_steps} steps.")
    else:
        # 如果距离较小，直接规划到目标位姿
        move_group.set_pose_target(pose.pose, link_name)
        plan = move_group.plan()

        if plan[0]:
            success = move_group.execute(plan[1], wait=True)
            if success:
                rospy.loginfo(f"Successfully moved {link_name} to the target pose.")
                return True
            else:
                rospy.logwarn(f"Failed to execute plan for {link_name}.")
                # 如果执行失败，尝试另外两种姿态
                return try_alternative_orientations(pose, link_name)
        else:
            rospy.logwarn(f"Failed to generate plan for {link_name}.")
            # 如果规划失败，尝试另外两种姿态
            return try_alternative_orientations(pose, link_name)
def try_alternative_orientations(pose, link_name):
    """
    尝试使用另外两种姿态进行规划。
    :param pose: 目标位姿 (Pose 或 PoseStamped)
    :param link_name: 要移动的链接名称
    :return: 是否成功
    """
    global move_group

    # 如果传入的是 PoseStamped 对象，提取其中的 Pose
    if isinstance(pose, PoseStamped):
        pose = pose.pose

    # 第一种备选姿态
    alternative_pose_1 = copy.deepcopy(pose)
    alternative_pose_1.orientation.x = -0.0321
    alternative_pose_1.orientation.y = 0.1140
    alternative_pose_1.orientation.z = 0.9236
    alternative_pose_1.orientation.w = 0.3645

    move_group.set_pose_target(alternative_pose_1, link_name)
    plan = move_group.plan()

    if plan[0]:
        success = move_group.execute(plan[1], wait=True)
        if success:
            rospy.loginfo(f"Successfully moved {link_name} to the target pose using alternative orientation 1.")
            return True
        else:
            rospy.logwarn(f"Failed to execute plan for {link_name} using alternative orientation 1.")
    else:
        rospy.logwarn(f"Failed to generate plan for {link_name} using alternative orientation 1.")

    # 第二种备选姿态
    alternative_pose_2 = copy.deepcopy(pose)
    alternative_pose_2.orientation.x = -0.014
    alternative_pose_2.orientation.y = 0.332
    alternative_pose_2.orientation.z = 0.912
    alternative_pose_2.orientation.w = 0.24

    move_group.set_pose_target(alternative_pose_2, link_name)
    plan = move_group.plan()

    if plan[0]:
        success = move_group.execute(plan[1], wait=True)
        if success:
            rospy.loginfo(f"Successfully moved {link_name} to the target pose using alternative orientation 2.")
            return True
        else:
            rospy.logwarn(f"Failed to execute plan for {link_name} using alternative orientation 2.")
    else:
        rospy.logwarn(f"Failed to generate plan for {link_name} using alternative orientation 2.")

    # 如果两种备选姿态都失败，则报错
    rospy.logwarn("Failed to move to target pose using all alternative orientations.")
    return False
def move_to_stand_by_pose():
    """
    移动手臂到初始位置。
    """
    global move_group
    stand_by_pose = PoseStamped()
    stand_by_pose.header.frame_id = "base_link"
    stand_by_pose.pose.position.x = 0.5
    stand_by_pose.pose.position.y = -0.5
    stand_by_pose.pose.position.z = 0.8
    stand_by_pose.pose.orientation.x = 0.11236370921869202
    stand_by_pose.pose.orientation.y = 0.15079563957693937
    stand_by_pose.pose.orientation.z = 0.7905900885524749
    stand_by_pose.pose.orientation.w = 0.5827541366801177

    return move_to_pose(stand_by_pose, "right_tool_link")


def adjust_pose_for_paper(pose):
    """
    调整目标位姿以适应 paper 任务。
    :param pose: 原始目标位姿 (PoseStamped)
    :return: 调整后的目标位姿 (PoseStamped)
    """
    adjusted_pose = copy.deepcopy(pose)
    # 调整位置
    adjusted_pose.pose.position.x += 0.04
    adjusted_pose.pose.position.y += -0.09
    adjusted_pose.pose.position.z += 0.15
    # 调整姿态
    adjusted_pose.pose.orientation.x = 0.11236370921869202
    adjusted_pose.pose.orientation.y = 0.15079563957693937
    adjusted_pose.pose.orientation.z = 0.7905900885524749
    adjusted_pose.pose.orientation.w = 0.5827541366801177
    return adjusted_pose


def task_manager():
    """
    任务管理器，根据任务状态执行相应的动作。
    """
    global task_state, start_move, left_pose, hand_control_publisher

    rate = rospy.Rate(10)  # 10Hz
    while not rospy.is_shutdown():
        if not start_move:
            rate.sleep()
            continue

        if task_state == "stand_by":
            rospy.loginfo("Task state: stand_by")
            
            # 检查机械臂是否在初始位置
            current_pose = move_group.get_current_pose().pose
            stand_by_pose = PoseStamped()
            stand_by_pose.header.frame_id = "base_link"
            stand_by_pose.pose.position.x = 0.5
            stand_by_pose.pose.position.y = -0.5
            stand_by_pose.pose.position.z = 0.8
            stand_by_pose.pose.orientation.x = 0.11236370921869202
            stand_by_pose.pose.orientation.y = 0.15079563957693937
            stand_by_pose.pose.orientation.z = 0.7905900885524749
            stand_by_pose.pose.orientation.w = 0.5827541366801177

            # 计算当前位置与初始位置的差距
            distance = np.linalg.norm(
                np.array([current_pose.position.x, current_pose.position.y, current_pose.position.z]) -
                np.array([stand_by_pose.pose.position.x, stand_by_pose.pose.position.y, stand_by_pose.pose.position.z])
            )
            orientation_diff = np.linalg.norm(
                np.array([current_pose.orientation.x, current_pose.orientation.y, current_pose.orientation.z, current_pose.orientation.w]) -
                np.array([stand_by_pose.pose.orientation.x, stand_by_pose.pose.orientation.y, stand_by_pose.pose.orientation.z, stand_by_pose.pose.orientation.w])
            )

            # 如果不在初始位置，则移动到初始位置
            #if distance > 0.1 or orientation_diff > 0.1:  # 位置和姿态的容忍误差
            if distance > 0.1:
                rospy.loginfo("Arm is not in stand_by pose, moving to stand_by pose...")
                if move_to_stand_by_pose():
                    rospy.loginfo("Arm is now in stand_by pose.")
                else:
                    rospy.logwarn("Failed to move to stand_by pose, retrying...")
                    continue


            # 等待目标位姿
            if left_pose is not None:
                task_state = "move_to_paper"
            else:
                print("left_pose",left_pose)
                rospy.loginfo("Waiting for target pose...")

        elif task_state == "move_to_paper":
            rospy.loginfo("Task state: move_to_paper")
            # 调整目标位姿
            adjusted_pose = adjust_pose_for_paper(left_pose)
            # 移动到调整后的 paper 位置
            if move_to_pose(adjusted_pose, "right_tool_link"):
                task_state = "grasp"
            else:
                rospy.logwarn("Failed to move to paper pose, retrying...")

        elif task_state == "grasp":
            rospy.loginfo("Task state: grasp")
            # 发送控制手的命令，执行关闭操作
            msg_hand_control = Bool()
            msg_hand_control.data = False  # False 表示关闭手
            hand_control_publisher.publish(msg_hand_control)
            rospy.loginfo("Sent hand close command.")
            rospy.sleep(2)  # 等待手完成动作
            task_state = "move_to_release"

        elif task_state == "move_to_release":
            rospy.loginfo("Task state: move_to_release")
            # 计算 release 位置
            release_pose = copy.deepcopy(left_pose)
            release_pose = adjust_pose_for_paper(release_pose)
            release_pose.pose.position.y += 0.1  # 假设 release 位置在 paper 位置的侧
            # 移动到 release 位置
            if move_to_pose(release_pose, "right_tool_link"):
                task_state = "release"
            else:
                rospy.logwarn("Failed to move to release pose, retrying...")

        elif task_state == "release":
            rospy.loginfo("Task state: release")
            # 发送控制手的命令，执行打开操作
            msg_hand_control = Bool()
            msg_hand_control.data = True  # True 表示打开手
            hand_control_publisher.publish(msg_hand_control)
            rospy.loginfo("Sent hand open command.")
            rospy.sleep(2)  # 等待手完成动作
            task_state = "return_home"

        elif task_state == "return_home":
            rospy.loginfo("Task state: return_home")
            # 返回初始位置
            if move_to_stand_by_pose():
                rospy.loginfo("Task completed, returning to stand_by state.")
                task_state = "stand_by"
                start_move = False  # 任务完成，等待下一次开始命令
            else:
                rospy.logwarn("Failed to return to stand_by pose, retrying...")

        rate.sleep()


if __name__ == "__main__":
    # 初始化 ROS 节点
    rospy.init_node('move_arm_task_manager', anonymous=True)

    # 初始化 MoveIt
    moveit_commander.roscpp_initialize(sys.argv)
    move_group = MoveGroupCommander("right_arm")
    move_group.set_end_effector_link("right_tool_link")
    move_group.set_planning_time(0.2)  # 增加规划时间
    move_group.set_num_planning_attempts(10)  # 增加尝试次数
    move_group.set_max_velocity_scaling_factor(0.3)
    move_group.set_max_acceleration_scaling_factor(0.3)
    move_group.set_goal_position_tolerance(0.01)  # 位置容差
    move_group.set_goal_orientation_tolerance(0.3)  # 姿态容差
    move_group.set_planner_id("PRM")  # 使用RRTConnect算法

    # 初始化手控制发布器
    hand_control_publisher = rospy.Publisher("/hand_control_topic", Bool, queue_size=1)

    # 订阅目标位姿和开始命令
    rospy.Subscriber('/paper_pose', PoseStamped, left_pose_callback, queue_size=1)
    rospy.Subscriber('/start_move', Bool, start_move_callback, queue_size=1)

    # 启动任务管理器
    task_manager()

    # 保持节点运行
    rospy.spin()